package cn.zengdeai.langchain4j_springboot.controller;

import cn.zengdeai.langchain4j_springboot.config.AiConfig;
import cn.zengdeai.langchain4j_springboot.tool.LogExample;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.service.TokenStream;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.text.Format;
import java.time.LocalDate;

@RestController
@RequestMapping("/ai")
public class ChatController {

    @Autowired
    QwenChatModel qwenChatModel;

    @Autowired
    QwenStreamingChatModel qwenStreamingChatModel;

    @Autowired
    AiConfig.Assistant assistant;

    @Autowired
    AiConfig.AssistantForCSR assistantForCSR;


    // 使用 SLF4J 获取一个 Logger 实例
    private static final Logger logger = LoggerFactory.getLogger(ChatController.class);

    /**
     * 非流式对话
     * @param message
     * @return
     */
    @RequestMapping("/chat")
    public String test(@RequestParam(defaultValue = "你是谁") String message) {
        //记录接口开始时间
        long startTime = System.currentTimeMillis();
        logger.info("request body is {},开始时间为:{}",message,formatTimestamp(startTime));
        String chat = qwenChatModel.chat(message);
        // 记录接口结束时间
        long endTime = System.currentTimeMillis();
        logger.info("response body is :{},结束时间为:{},接口总耗时为:{} 毫秒",chat, formatTimestamp(System.currentTimeMillis()),endTime-startTime);
        return chat;
    }


    /**
     * 流式对话准备工作，未生效
     * @param message
     * @return
     */
    @RequestMapping("/stream_ori")
    public String streamChatOri(@RequestParam(defaultValue = "你是谁") String message) {
        qwenStreamingChatModel.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String s) {
                System.out.println(s);
            }

            @Override
            public void onCompleteResponse(ChatResponse chatResponse) {
                System.out.println("已经结束");
            }

            @Override
            public void onError(Throwable throwable) {
                throwable.printStackTrace();
            }
        });
        return "chat";
    }


    /**
     * 流式对话
     * @param message
     * @return
     */
    @RequestMapping(value = "/stream", produces = "text/stream;charset=UTF-8")
    public Flux<String> streamChat(@RequestParam(defaultValue = "你是谁") String message) {

        Flux<String> objectFlux = Flux.create(fluxSink -> {
            qwenStreamingChatModel.chat(message, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    fluxSink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    fluxSink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    fluxSink.error(throwable);
                }
            });

        });
        return objectFlux;
    }

    /**
     * 多轮对话记忆+非流式
     * @param message
     * @return
     */
    @RequestMapping(value = "/memory_chat")
    public String memoryChat(@RequestParam(defaultValue = "我叫万德军") String message) {
        return assistant.chat(message);
    }


    /**
     * 多轮对话记忆+流式
     * @param message
     * @param response
     * @return
     */
    @RequestMapping(value = "/memory_stream_chat", produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryStreamChat(@RequestParam(defaultValue = "我是谁") String message, HttpServletResponse response) {
        TokenStream stream = assistant.stream(message);

        return Flux.create(sink -> {
            stream.onPartialResponse(s -> sink.next(s))
                    .onCompleteResponse(c -> sink.complete())
                    .onError(sink::error)
                    .start();

        });
    }


    /**
     * 多轮对话记忆+流式
     * @param message
     * @param response
     * @return
     */
    @RequestMapping(value = "/system_message_chat", produces = "text/stream;charset=UTF-8")
    public Flux<String> systemMessageChat(@RequestParam(defaultValue = "我是谁") String message, HttpServletResponse response) {
        TokenStream stream = assistant.stream(message, LocalDate.now().toString());

        return Flux.create(sink -> {
            stream.onPartialResponse(s -> sink.next(s))
                    .onCompleteResponse(c -> sink.complete())
                    .onError(sink::error)
                    .start();

        });
    }

    /**
     * 多轮对话记忆流式:训练客服
     * @param message
     * @param response
     * @return
     */
    @RequestMapping(value = "/stream_for_train_for_csr", produces = "text/stream;charset=UTF-8")
    public Flux<String> systemMessageForCSR(@RequestParam(defaultValue = "我是谁") String message, HttpServletResponse response) {
        TokenStream stream = assistantForCSR.streamForTrainForCSR(message, LocalDate.now().toString());

        return Flux.create(sink -> {
            stream.onPartialResponse(s -> sink.next(s))
                    .onCompleteResponse(c -> sink.complete())
                    .onError(sink::error)
                    .start();

        });
    }







    @Autowired
    AiConfig.AssistantUnique assistantUnique;

    /**
     * 多租户隔离
     * @param message
     * @param userId
     * @return
     */
    @RequestMapping(value = "/memoryId_chat")
    public String memoryChat(@RequestParam(defaultValue="我是谁") String message, Integer userId) {
        return assistantUnique.chat(userId,message);
    }



    @RequestMapping(value = "/function_call_chat")
    public String functionCallChat(@RequestParam(defaultValue="我是谁") String message) {
//        return  assistant.stream(message);
        return assistant.chat(message);
    }


    // 格式化时间戳为可读格式
    private String formatTimestamp(long timestamp) {
        return new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS").format(new java.util.Date(timestamp));
    }

}
